import re, json, os, sys
import time
import traceback
import requests
import subprocess
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QIcon
from PyQt5.QtWidgets import (QApplication, QWidget, QVBoxLayout, QLineEdit, QPushButton, QTextEdit, QProgressBar, QMessageBox, QLabel, QHBoxLayout, QFileDialog)
from pathlib import Path
from PyQt5.QtCore import QThread, pyqtSignal
import crypto
from crypto import decrypt_from_txt, encrypt_to_txt, load_key  # 修改导入

HEADERS = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
    "Referer": "",
}

def extract_title(html: str) -> str:
    m = re.search(r'<title>([^<]+)</title>', html)
    return m.group(1).split('_')[0].strip() if m else "unknown"

def safe_name(title: str, ext: str, max_len=80) -> str:
    invalid_chars = r'[\\/:*?"<>|！？【】{}「」『』《》]'
    title = re.sub(invalid_chars, '', title).replace(' ', '_')
    return (title[:max_len] if len(title) > max_len else title) + f".{ext}"

def unique_path(dir_: Path, name: str) -> Path:
    """同名自动加 _1 _2..."""
    candidate = dir_ / name
    if not candidate.exists():
        return candidate
    idx = 1
    stem, suffix = Path(name).stem, Path(name).suffix
    while True:
        candidate = dir_ / f"{stem}_{idx}{suffix}"
        if not candidate.exists():
            return candidate

class DownloadTask(QThread):
    progress = pyqtSignal(int)
    log      = pyqtSignal(str)
    finished = pyqtSignal(bool)
    paths    = pyqtSignal(object)   # dict

    def __init__(self, url, cookie, download_dir: Path):
        super().__init__()
        self.url, self.cookie, self.download_dir = url, cookie, download_dir
        headers = HEADERS.copy()
        headers["Referer"] = url
        headers["Cookie"]  = cookie
        self.html = requests.get(url, headers=headers, timeout=15).text

    def run(self):
        try:
            title = extract_title(self.html)
            dir_ = self.download_dir
            dir_.mkdir(parents=True, exist_ok=True)

            video_name = safe_name(title, "mp4")
            audio_name = safe_name(title, "mp3")
            final_name = safe_name(title, "mp4")

            self.video_path = unique_path(dir_, video_name)
            self.audio_path = unique_path(dir_, audio_name)
            self.final_path = unique_path(dir_, final_name)

            self.paths.emit({"video": str(self.video_path),
                             "audio": str(self.audio_path),
                             "final": str(self.final_path)})

            info = re.search(r'window\.__playinfo__=({.*?})</script>', self.html).group(1)
            data = json.loads(info)
            v_url = data['data']['dash']['video'][0]['baseUrl']
            a_url = data['data']['dash']['audio'][0]['baseUrl']

            self.log.emit("开始下载视频…")
            self.download_stream(v_url, self.video_path)
            self.log.emit("开始下载音频…")
            self.download_stream(a_url, self.audio_path)
            self.finished.emit(True)
        except Exception as e:
            self.log.emit(f"失败：{e}")
            self.finished.emit(False)

    def download_stream(self, url, file_path: Path):
        headers = HEADERS.copy()
        headers["Referer"] = self.url
        headers["Cookie"] = self.cookie
        total = int(requests.head(url, headers=headers).headers.get('content-length', 0))
        done = 0
        last_percent = 0
        self.log.emit(f"开始写入 {file_path.name} …")

        with requests.get(url, headers=headers, stream=True) as r:
            for chunk in r.iter_content(1024 * 64):
                if chunk:
                    with open(file_path, 'ab') as f:
                        f.write(chunk)
                    done += len(chunk)
                    if total:
                        percent = int(done * 100 / total)
                    else:
                        percent = int(done / 1024 / 1024)
                    if percent != last_percent:
                        self.progress.emit(percent)
                        last_percent = percent
        self.progress.emit(100)

class AudioOnlyTask(QThread):
    progress = pyqtSignal(int)
    log = pyqtSignal(str)
    finished = pyqtSignal(bool)
    path = pyqtSignal(str)          # 把最终 mp3 路径抛出去

    def __init__(self, url, cookie, download_dir: Path):
        super().__init__()
        self.url, self.cookie, self.download_dir = url, cookie, download_dir
        headers = HEADERS.copy()
        headers["Referer"] = url
        headers["Cookie"]  = cookie
        self.html = requests.get(url, headers=headers, timeout=15).text

    def run(self):
        try:
            title = extract_title(self.html)
            dir_ = Path("download")
            dir_.mkdir(exist_ok=True)
            audio_name = safe_name(title, "mp3")
            self.audio_path = unique_path(dir_, audio_name)
            self.path.emit(str(self.audio_path))

            info = re.search(r'window\.__playinfo__=({.*?})</script>', self.html).group(1)
            data = json.loads(info)
            a_url = data['data']['dash']['audio'][0]['baseUrl']

            self.log.emit("开始下载音频…")
            self.download_stream(a_url, self.audio_path)
            self.finished.emit(True)
        except Exception as e:
            self.log.emit(f"音频下载失败：{e}")
            self.finished.emit(False)

    def download_stream(self, url, file_path: Path):
        headers = HEADERS.copy()
        headers["Referer"] = self.url
        headers["Cookie"] = self.cookie
        total = int(requests.head(url, headers=headers).headers.get('content-length', 0))
        done = last_percent = 0
        self.log.emit(f"开始写入 {file_path.name} …")
        with requests.get(url, headers=headers, stream=True) as r:
            for chunk in r.iter_content(1024 * 64):
                if chunk:
                    with open(file_path, 'ab') as f:
                        f.write(chunk)
                    done += len(chunk)
                    percent = int(done * 100 / total) if total else int(done / 1024 / 1024)
                    if percent != last_percent:
                        self.progress.emit(percent)
                        last_percent = percent
        self.progress.emit(100)

class MergeTask(QThread):
    progress = pyqtSignal(int)   # 0-100
    log      = pyqtSignal(str)
    finished = pyqtSignal(bool)

    def __init__(self, video, audio, out):
        super().__init__()
        self.video = str(video)
        self.audio = str(audio)
        self.out = str(Path(out).with_name(Path(out).stem + "_complete.mp4"))

    def run(self):
        try:
            self.log.emit("正在合并音视频…")
            cmd = [
                "ffmpeg", "-i", self.video,
                          "-i", self.audio,
                          "-c:v", "copy", "-c:a", "aac", "-strict", "experimental",
                          "-progress", "pipe:1",   # 把进度输出到管道
                          "-y", self.out
            ]
            total_size = os.path.getsize(self.video) + os.path.getsize(self.audio)
            done = 0
            with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                  bufsize=1, universal_newlines=True, encoding='utf-8') as p:
                last_log = 0
                for line in p.stdout:
                    if "out_time_ms=" in line:
                        ms = int(line.split("=")[1])
                        done = min(ms // 1000, 100)
                        self.progress.emit(done)
                    if time.time() - last_log > 0.5:
                        self.log.emit(line.rstrip())
                        last_log = time.time()
            if p.returncode == 0:
                self.finished.emit(True)
            else:
                self.finished.emit(False)
        except Exception as e:
            self.log.emit(str(e))
            self.finished.emit(False)
            self.log.emit(f"[MergeTask] 异常：{traceback.format_exc()}")

def merge(video, audio, out="output.mp4"):
    cmd = [
        "ffmpeg", "-i", video,
                    "-i", audio,
                    "-c:v", "copy",
                    "-c:a", "aac", "-strict", "experimental",
                    "-y", out
    ]
    subprocess.run(cmd, check=True)
    print(f"合并完成 → {out}")

# 加载密钥
def load_key():
    try:
        return crypto.load_key(Path("secret.key"))
    except FileNotFoundError:
        return None


def load_cookie():
    try:
        key = load_key()
        if key is None:
            print("密钥文件不存在")
            return ""
        print(f"密钥加载成功: {key[:20]}... (长度: {len(key)})")

        decrypted = crypto.decrypt_from_txt(Path("cookie.bin"), key)
        print(f"Cookie 解密成功 (长度: {len(decrypted)})")
        return decrypted
    except Exception as e:
        print(f"load_cookie 异常: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()  # 打印完整堆栈
        return ""

def load_download_dir():
    key = load_key()
    if key is None:
        return Path("download")
    try:
        return Path(crypto.decrypt_from_txt(Path("download_dir.bin"), key))
    except FileNotFoundError:
        return Path("download")

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("B 站下载器  |  PyQt5 版")
        self.resize(1440, 960)
        self.cookie = load_cookie()
        if not self.cookie:
            print("警告: Cookie 为空！请检查 secret.key 和 cookie.bin 是否存在")

        ico_path = Path(__file__).with_name("app.ico").absolute()
        if ico_path.exists():
            self.setWindowIcon(QIcon(str(ico_path)))
        else:
            print("图标文件不存在:", ico_path)

        # 1. 控件
        self.url_edit = QLineEdit()
        self.url_edit.setPlaceholderText("粘贴 B 站视频地址，回车下载")
        self.btn_layout = QHBoxLayout()
        self.download_btn = QPushButton("下载")
        self.audio_only_btn = QPushButton("仅下载音频")
        self.dir_btn = QPushButton("选择下载目录")
        self.btn_layout.addWidget(self.download_btn)
        self.btn_layout.addWidget(self.audio_only_btn)
        self.btn_layout.addWidget(self.dir_btn)

        self.status = QLabel("就绪")
        self.status.setAlignment(Qt.AlignCenter)
        self.pb = QProgressBar()
        self.log = QTextEdit(readOnly=True)

        # 2. 布局
        lay = QVBoxLayout(self)
        lay.addWidget(self.url_edit)
        lay.addLayout(self.btn_layout)
        lay.addWidget(self.status)
        lay.addWidget(self.pb)
        lay.addWidget(self.log)

        # 3. 信号
        self.download_btn.clicked.connect(self.start_download)
        self.url_edit.returnPressed.connect(self.start_download)
        self.audio_only_btn.clicked.connect(self.start_audio_only)
        self.dir_btn.clicked.connect(self.choose_download_dir)

        print('cookie 前 100 字符：', self.cookie[:100])

    def start_download(self):
        url = self.url_edit.text().strip()
        if not url:
            QMessageBox.warning(self, "提示", "地址不能为空")
            return

        # 防止重复启动任务
        if hasattr(self, 'task') and self.task and self.task.isRunning():
            QMessageBox.warning(self, "提示", "下载任务正在进行中...")
            return

        self.pb.setValue(0)
        self.log.clear()  # 清空旧日志
        self.log.append("任务开始……")

        try:
            self.task = DownloadTask(url, self.cookie, self.load_download_dir())
            # 使用 QueuedConnection 确保线程安全
            self.task.progress.connect(self.pb.setValue, Qt.QueuedConnection)
            self.task.log.connect(self.log.append, Qt.QueuedConnection)
            self.task.finished.connect(self.done, Qt.QueuedConnection)
            self.task.start()
        except Exception as e:
            self.log.append(f"启动任务失败: {e}")
            QMessageBox.critical(self, "错误", f"无法启动下载任务:\n{e}")

    def done(self, ok: bool):
        if not ok:
            QMessageBox.warning(self, "失败", "下载未完成，无法合并")
            self.status.setText("就绪")
            return

        self.status.setText("正在合并…")
        self.log.append(">>> 开始合并音视频")
        v_file = self.task.video_path
        a_file = self.task.audio_path
        c_file = self.task.final_path

        self.merge_task = MergeTask(v_file, a_file, c_file)
        self.merge_task.progress.connect(self.pb.setValue)
        self.merge_task.log.connect(self.log.append)
        self.merge_task.finished.connect(lambda succ: self.merge_done(succ, v_file, a_file, c_file))
        self.merge_task.start()

    def merge_done(self, succ, v_file, a_file, c_file):
        if succ:
            QMessageBox.information(self, "完成", f"合并成功！\n{c_file}")
            Path(v_file).unlink(missing_ok=True)
            Path(a_file).unlink(missing_ok=True)
            self.status.setText("合并完成")
        else:
            QMessageBox.warning(self, "合并失败", "请看日志")
            self.status.setText("就绪")

    def start_audio_only(self):
        url = self.url_edit.text().strip()
        if not url:
            QMessageBox.warning(self, "提示", "地址不能为空")
            return
        self.pb.setValue(0)
        self.log.append("仅下载音频任务开始……")
        self.audio_task = AudioOnlyTask(url, self.cookie, self.load_download_dir())
        self.audio_task.progress.connect(self.pb.setValue)
        self.audio_task.log.connect(self.log.append)
        self.audio_task.finished.connect(self.audio_done)
        self.audio_task.start()

    def audio_done(self, ok: bool):
        if ok:
            mp3 = self.audio_task.audio_path
            QMessageBox.information(self, "完成", f"音频已保存为\n{mp3}")
            self.status.setText("音频下载完成")
        else:
            QMessageBox.warning(self, "失败", "音频下载失败，请看日志")
            self.status.setText("就绪")

    def load_download_dir(self):
        key = load_key()
        if key is None:
            return Path("download")
        try:
            return Path(decrypt_from_txt(Path("download_dir.bin"), key))
        except FileNotFoundError:
            return Path("download")

    def choose_download_dir(self):
        key = load_key()
        if key is None:
            return
        old = self.load_download_dir()
        new = QFileDialog.getExistingDirectory(
            self, "选择视频保存文件夹", str(old)
        )
        if new:
            new_path = Path(new)
            encrypt_to_txt(str(new_path), Path("download_dir.bin"), key)
            QMessageBox.information(self, "提示", f"下载目录已设置为：\n{new_path}")

if __name__ == '__main__':
    app = QApplication(sys.argv)
    w = MainWindow()
    w.show()
    sys.exit(app.exec())